# -*- coding: utf-8 -*-
"""
@author: wyb
@date: 2024/1/12
@time: 20:45
@desc: 萝卜投研 API封装
"""
import os.path
from urllib.parse import urljoin
from selenium import webdriver
from selenium.webdriver.common.by import By
import requests
import time
import json, csv


def writeStockToCSV(data, filename):
    '''
     将股票数据写入csv文件
    :param data: 股票数据，包括代码和名称
    :param filename: 文件名称
    :return:
    '''
    with open(filename, 'w', newline='') as csvfile:
        fieldnames = ['ticker', 'name']  # CSV文件的列名
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()  # 写入列名作为表头
        for row in data:
            writer.writerow(row)  # 写入每行数据


def loginrobo(url, username, pwd):
    '''
     登录 萝卜投研
    :param url: 登录地址
    :param username: 用户名
    :param pwd: 密码
    :return: 登陆后填充cookie后的会话
    '''
    browser = webdriver.Chrome()
    browser.get(url)
    time.sleep(5)
    switchEle = browser.find_element(By.CSS_SELECTOR, 'div.login-switch-btn > img')
    switchEle.click()
    time.sleep(3)
    userEle = browser.find_element(By.CSS_SELECTOR, 'input[id="username"]')
    userEle.send_keys(username)

    pwdEle = browser.find_element(By.CSS_SELECTOR, 'input[id="password"]')
    pwdEle.send_keys(pwd)

    rememberEle = browser.find_element(By.CSS_SELECTOR, 'input[id="rememberMe"]')
    rememberEle.click()

    submitEle = browser.find_element(By.CSS_SELECTOR, 'button[type="submit"]')
    submitEle.click()
    time.sleep(5)

    # get cookies from selenium
    cookies = browser.get_cookies()
    browser.close()

    # set cookies to requests
    session = requests.Session()
    for cookie in cookies:
        session.cookies.set(cookie['name'], cookie['value'])
    return session


def downFinanceReport(session, lis, subdir):
    '''
     下载财务报告为本地文件 excel 文件
     https://gw.datayes.com/rrp_adventure/fdmtNew/excel/000895?decimalPoint=2&displaySort=left&duration=ACCUMULATE&includeLatest=true&isHideEmptyData=true&period=10&reportPeriodType=A&reportType=BS&ticker=000895&unit=2
     https://gw.datayes.com/rrp_adventure/fdmtNew/excel/000895?decimalPoint=2&displaySort=left&duration=ACCUMULATE&includeLatest=true&isHideEmptyData=true&period=10&reportPeriodType=A&reportType=IS&ticker=000895&unit=2
     https://gw.datayes.com/rrp_adventure/fdmtNew/excel/000895?decimalPoint=2&displaySort=left&duration=ACCUMULATE&includeLatest=true&isHideEmptyData=true&period=10&reportPeriodType=A&reportType=CF&ticker=000895&unit=2
    :param session: 会话
    :param url: 请求地址
    :param filename: 下载后的文件名
    :return:
    '''
    reporttype = ['BS', 'IS', 'CF']
    for stock in lis:
        for t in reporttype:
            time.sleep(1)
            url = f'https://gw.datayes.com/rrp_adventure/fdmtNew/excel/{stock}?decimalPoint=2&displaySort=left&duration=ACCUMULATE&includeLatest=true&isHideEmptyData=true&period=10&reportPeriodType=A&reportType={t}&ticker={stock}&unit=2'
            response = session.get(url)
            filename = os.path.join(subdir, stock + '_' + t + ".xls")
            with open(filename, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)


def choseStocks(session):
    '''
     根据条件选股票
    :param session: 请求会话
    :return: 股票名称和代码列表
    '''
    postStockList = 'https://gw.datayes.com/rrp_adventure/stockfactor/stocks'
    postData = [{"tab": "常用指标", "category": "", "factorName": "总市值", "factorDisplayName": "总市值(元)",
                 "factorValue": "[100亿, ]", "uniqId": 53, "unit": "亿", "firstLetter": "zd",
                 "customValue": "[100亿, ]"},
                {"tab": "股票范围", "category": "市场板块", "factorName": "", "factorDisplayName": "",
                 "factorValue": "上证A股", "uniqId": 1, "unit": "", "firstLetter": "szA", "customValue": 'null'},
                {"tab": "常用指标", "category": "", "factorName": "市盈率", "factorDisplayName": "市盈率(倍)",
                 "factorValue": "10~20", "uniqId": 316, "unit": "", "firstLetter": "1", "customValue": 'null'}]
    headers = {
        'Content-Type': 'application/json;charset=UTF-8'  # 通常设置为application/json，但根据实际情况可能有所不同
    }

    response = session.post(postStockList, json=postData, headers=headers)
    try:
        response.raise_for_status()
        # HTTP 请求成功，继续处理响应数据
        if response.status_code == 200:
            result = response.text
            data = json.loads(result)
            lis = data['data']['stockList']
            stocks = [{'ticker': obj['ticker'], 'name': obj['secShortName']} for obj in lis]
            # json_data = json.loads(stocks)
            return stocks
            writeStockToCSV(stocks)
            print(stocks)
        else:
            print(f"Error: {response.status_code}")
    except requests.HTTPError as e:
        print(f"HTTP请求错误: {e}")
